import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import os
import cv2
import threading
from PIL import Image, ImageTk, ImageDraw
import re
from concurrent.futures import ThreadPoolExecutor


class VideoFrameExtractor:
    def __init__(self, root):
        self.root = root
        self.root.title("智能视频截图工具")
        self.root.geometry("800x700")
        self.root.configure(bg="#f0f2f5")

        # 初始化变量
        self.video_path = ""
        self.output_folder = ""
        self.interval_seconds = 1.0
        self.processing = False
        self.cancel_flag = False
        self.max_workers = 2  # 限制线程数避免FFmpeg冲突
        self.cap = None
        self.preview_cap = None

        # 创建界面元素
        self.create_widgets()

    def sanitize_filename(self, name):
        """清理文件名中的非法字符"""
        return re.sub(r'[\\/*?:"<>|]', "_", name)

    def create_widgets(self):
        # 主容器
        main_frame = tk.Frame(self.root, bg="#f0f2f5", padx=20, pady=20)
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 标题
        title_frame = tk.Frame(main_frame, bg="#f0f2f5")
        title_frame.pack(fill=tk.X, pady=(0, 10))

        tk.Label(title_frame, text="智能视频截图工具", font=("Arial", 18, "bold"),
                 bg="#f0f2f5", fg="#2c3e50").pack(side=tk.LEFT)

        # 文件选择区域
        file_frame = tk.LabelFrame(main_frame, text="文件选择", font=("Arial", 12),
                                   bg="#ffffff", fg="#2c3e50", padx=15, pady=15)
        file_frame.pack(fill=tk.X, pady=5)

        # 视频选择
        video_frame = tk.Frame(file_frame, bg="#ffffff")
        video_frame.pack(fill=tk.X, pady=5)

        tk.Label(video_frame, text="视频文件:", font=("Arial", 11),
                 bg="#ffffff", fg="#34495e").pack(side=tk.LEFT, padx=(0, 10))

        self.video_path_entry = ttk.Entry(video_frame, font=("Arial", 11))
        self.video_path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 5))

        ttk.Button(video_frame, text="浏览", command=self.select_video).pack(side=tk.LEFT)

        # 输出文件夹选择
        output_frame = tk.Frame(file_frame, bg="#ffffff")
        output_frame.pack(fill=tk.X, pady=5)

        tk.Label(output_frame, text="输出文件夹:", font=("Arial", 11),
                 bg="#ffffff", fg="#34495e").pack(side=tk.LEFT, padx=(0, 10))

        self.output_path_entry = ttk.Entry(output_frame, font=("Arial", 11))
        self.output_path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 5))

        ttk.Button(output_frame, text="浏览", command=self.select_output_folder).pack(side=tk.LEFT)

        # 设置区域
        setting_frame = tk.LabelFrame(main_frame, text="设置", font=("Arial", 12),
                                      bg="#ffffff", fg="#2c3e50", padx=15, pady=15)
        setting_frame.pack(fill=tk.X, pady=5)

        # 间隔秒数输入
        interval_frame = tk.Frame(setting_frame, bg="#ffffff")
        interval_frame.pack(fill=tk.X, pady=5)

        tk.Label(interval_frame, text="截图间隔秒数:", font=("Arial", 11),
                 bg="#ffffff", fg="#34495e").pack(side=tk.LEFT, padx=(0, 10))

        self.interval_entry = ttk.Entry(interval_frame, font=("Arial", 11), width=10)
        self.interval_entry.insert(0, "1.0")
        self.interval_entry.pack(side=tk.LEFT, padx=(0, 5))

        tk.Label(interval_frame, text="秒", font=("Arial", 11),
                 bg="#ffffff", fg="#34495e").pack(side=tk.LEFT)

        # 文件名格式设置
        filename_frame = tk.Frame(setting_frame, bg="#ffffff")
        filename_frame.pack(fill=tk.X, pady=5)

        tk.Label(filename_frame, text="文件名格式:", font=("Arial", 11),
                 bg="#ffffff", fg="#34495e").pack(side=tk.LEFT, padx=(0, 10))

        # 前缀输入
        self.prefix_entry = ttk.Entry(filename_frame, font=("Arial", 11), width=10)
        self.prefix_entry.insert(0, "frame")
        self.prefix_entry.pack(side=tk.LEFT)

        # 数字格式
        num_frame = tk.Frame(filename_frame, bg="#ffffff")
        num_frame.pack(side=tk.LEFT, padx=5)

        tk.Label(num_frame, text="+", font=("Arial", 11),
                 bg="#ffffff", fg="#7f8c8d").pack(side=tk.LEFT)

        self.digits_entry = ttk.Entry(num_frame, font=("Arial", 11), width=5)
        self.digits_entry.insert(0, "0000")
        self.digits_entry.pack(side=tk.LEFT)

        tk.Label(num_frame, text="(编号)", font=("Arial", 10),
                 bg="#ffffff", fg="#7f8c8d").pack(side=tk.LEFT)

        # 后缀输入
        self.suffix_entry = ttk.Entry(filename_frame, font=("Arial", 11), width=10)
        self.suffix_entry.pack(side=tk.LEFT, padx=5)

        tk.Label(filename_frame, text=".jpg", font=("Arial", 11),
                 bg="#ffffff", fg="#7f8c8d").pack(side=tk.LEFT)

        # 控制按钮区域
        control_frame = tk.Frame(main_frame, bg="#f0f2f5")
        control_frame.pack(fill=tk.X, pady=5)

        # 进度条
        self.progress = ttk.Progressbar(control_frame, orient=tk.HORIZONTAL,
                                        length=400, mode='determinate')
        self.progress.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 10))

        # 开始/取消按钮
        self.start_button = ttk.Button(control_frame, text="开始截图",
                                       command=self.toggle_extraction)
        self.start_button.pack(side=tk.RIGHT)

        # 状态信息
        self.status_var = tk.StringVar()
        self.status_var.set("准备就绪")
        self.status_label = tk.Label(control_frame, textvariable=self.status_var,
                                     font=("Arial", 10), bg="#f0f2f5", fg="#2c3e50")
        self.status_label.pack(side=tk.LEFT, padx=(0, 10))

        # 预览区域
        preview_frame = tk.LabelFrame(main_frame, text="视频预览", font=("Arial", 12),
                                      bg="#ffffff", fg="#2c3e50", padx=15, pady=15)
        preview_frame.pack(fill=tk.BOTH, expand=True, pady=5)

        self.preview_label = tk.Label(preview_frame, bg="#2c3e50", fg="white",
                                      text="视频预览将显示在这里", font=("Arial", 11))
        self.preview_label.pack(fill=tk.BOTH, expand=True)

        # 配置样式
        self.configure_styles()

    def configure_styles(self):
        """配置ttk样式"""
        style = ttk.Style()
        style.theme_use('clam')

        # 按钮样式
        style.configure('TButton', font=('Arial', 11), padding=6)

        # 进度条样式
        style.configure("Horizontal.TProgressbar", thickness=20,
                        background='#3498db', troughcolor='#ecf0f1')

    def select_video(self):
        file_path = filedialog.askopenfilename(
            title="选择视频文件",
            filetypes=[("视频文件", "*.mp4 *.avi *.mov *.mkv *.flv *.wmv"), ("所有文件", "*.*")]
        )
        if file_path:
            self.video_path = file_path
            self.video_path_entry.delete(0, tk.END)
            self.video_path_entry.insert(0, file_path)
            self.update_preview()

    def select_output_folder(self):
        folder_path = filedialog.askdirectory(title="选择输出文件夹")
        if folder_path:
            self.output_folder = folder_path
            self.output_path_entry.delete(0, tk.END)
            self.output_path_entry.insert(0, folder_path)

    def update_preview(self):
        """更新视频预览"""
        if not self.video_path or not os.path.exists(self.video_path):
            return

        # 释放之前的预览捕获对象
        if self.preview_cap is not None:
            self.preview_cap.release()

        try:
            self.preview_cap = cv2.VideoCapture(self.video_path)
            if self.preview_cap.isOpened():
                ret, frame = self.preview_cap.read()
                if ret:
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    img = Image.fromarray(frame)
                    img.thumbnail((600, 300))
                    photo = ImageTk.PhotoImage(image=img)

                    self.preview_label.config(image=photo, text="")
                    self.preview_label.image = photo
                else:
                    self.preview_label.config(image=None, text="无法读取视频帧")
            else:
                self.preview_label.config(image=None, text="无法打开视频文件")
        except Exception as e:
            self.preview_label.config(image=None, text=f"预览错误: {str(e)}")

    def toggle_extraction(self):
        """切换截图状态"""
        if self.processing:
            self.cancel_flag = True
            self.status_var.set("正在停止...")
        else:
            self.start_extraction()

    def start_extraction(self):
        # 验证输入
        if not self.video_path:
            messagebox.showwarning("警告", "请先选择视频文件!")
            return

        if not self.output_folder:
            messagebox.showwarning("警告", "请先选择输出文件夹!")
            return

        try:
            self.interval_seconds = float(self.interval_entry.get())
            if self.interval_seconds <= 0:
                raise ValueError
        except ValueError:
            messagebox.showwarning("警告", "请输入有效的间隔秒数(正数)!")
            return

        # 检查文件名格式
        try:
            prefix = self.sanitize_filename(self.prefix_entry.get())
            digits = self.digits_entry.get()
            suffix = self.sanitize_filename(self.suffix_entry.get())
            test_name = f"{prefix}{digits}{suffix}.jpg"
            if not test_name.isprintable():
                raise ValueError("包含不可打印字符")
        except Exception as e:
            messagebox.showerror("错误", f"文件名格式无效: {str(e)}")
            return

        # 检查视频文件是否存在
        if not os.path.exists(self.video_path):
            messagebox.showerror("错误", "视频文件不存在!")
            return

        # 检查输出文件夹是否存在，不存在则创建
        if not os.path.exists(self.output_folder):
            try:
                os.makedirs(self.output_folder)
            except OSError as e:
                messagebox.showerror("错误", f"无法创建输出文件夹: {e}")
                return

        # 更新UI状态
        self.processing = True
        self.cancel_flag = False
        self.start_button.config(text="取消")
        self.status_var.set("正在处理...")
        self.progress['value'] = 0
        self.root.update()

        # 在新线程中开始处理
        threading.Thread(target=self.extract_frames, daemon=True).start()

    def extract_frames(self):
        try:
            # 创建新的视频捕获对象用于处理
            self.cap = cv2.VideoCapture(self.video_path)
            if not self.cap.isOpened():
                raise Exception("无法打开视频文件")

            fps = self.cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = total_frames / fps

            # 计算需要截取的帧时间点
            capture_times = []
            current_time = 0.0
            while current_time <= duration:
                capture_times.append(current_time)
                current_time += self.interval_seconds

            # 准备文件名参数
            prefix = self.sanitize_filename(self.prefix_entry.get())
            digits_format = self.digits_entry.get()
            suffix = self.sanitize_filename(self.suffix_entry.get())

            # 使用线程池处理（限制线程数）
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = []
                frame_count = 0
                total_captures = len(capture_times)

                for i, target_time in enumerate(capture_times):
                    if self.cancel_flag:
                        break

                    # 提交任务到线程池
                    future = executor.submit(
                        self.process_single_frame,
                        target_time, i, total_captures,
                        prefix, digits_format, suffix, frame_count
                    )
                    futures.append(future)
                    frame_count += 1

                # 等待所有任务完成
                for future in futures:
                    future.result()
                    if self.cancel_flag:
                        break

            if not self.cancel_flag:
                self.status_var.set("处理完成")
                messagebox.showinfo("完成", f"成功保存 {frame_count} 张截图")
        except Exception as e:
            if not self.cancel_flag:
                self.status_var.set(f"错误: {str(e)}")
                messagebox.showerror("错误", f"处理过程中出错: {str(e)}")
        finally:
            if self.cap is not None:
                self.cap.release()
            self.processing = False
            self.start_button.config(text="开始截图")
            if self.cancel_flag:
                self.status_var.set("已取消")
            self.root.update()

    def process_single_frame(self, target_time, i, total_captures,
                             prefix, digits_format, suffix, frame_count):
        """处理单个帧的多线程函数"""
        if self.cancel_flag:
            return

        try:
            # 创建独立的视频捕获对象
            thread_cap = cv2.VideoCapture(self.video_path)
            if not thread_cap.isOpened():
                raise Exception("无法打开视频文件")

            # 设置视频到目标时间点
            thread_cap.set(cv2.CAP_PROP_POS_MSEC, target_time * 1000)

            # 读取帧
            ret, frame = thread_cap.read()
            if ret:
                # 生成文件名
                num_str = f"{frame_count:{digits_format.replace('0', '')}}"
                output_path = os.path.join(
                    self.output_folder,
                    f"{prefix}{num_str}{suffix}.jpg"
                )
                cv2.imwrite(output_path, frame)

                # 更新进度 (需要在主线程执行)
                self.root.after(0, self.update_progress, i + 1, total_captures)
        finally:
            if thread_cap.isOpened():
                thread_cap.release()

    def update_progress(self, current, total):
        """更新进度条"""
        progress = int(current / total * 100)
        self.progress['value'] = progress
        self.status_var.set(f"处理中: {progress}% ({current}/{total})")

    def __del__(self):
        """析构函数确保释放资源"""
        if self.cap is not None and self.cap.isOpened():
            self.cap.release()
        if self.preview_cap is not None and self.preview_cap.isOpened():
            self.preview_cap.release()


if __name__ == "__main__":
    # 设置FFmpeg日志级别（抑制警告）
    os.environ['FFREPORT'] = 'file=ffreport.log:level=32'

    root = tk.Tk()
    app = VideoFrameExtractor(root)
    root.mainloop()